# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from torch import nn
from torch.nn import functional as F

from .roi_box_feature_extractors import make_roi_box_feature_extractor
from .roi_box_predictors import make_roi_box_predictor
from .inference import make_roi_box_post_processor
from .loss import make_roi_box_loss_evaluator
from maskrcnn_benchmark.modeling.utils import cat
from maskrcnn_benchmark.layers.smooth_l1_loss import smooth_l1_loss
from maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou
from maskrcnn_benchmark.structures.bounding_box import BoxList
class ROIBoxHead(torch.nn.Module):
    """
    Generic Box Head class.
    """

    def __init__(self, cfg, in_channels):
        super(ROIBoxHead, self).__init__()
        self.cfg = cfg
        self.feature_extractor = make_roi_box_feature_extractor(cfg, in_channels)
        self.predictor = make_roi_box_predictor(cfg, self.feature_extractor.out_channels)
        self.post_processor = make_roi_box_post_processor(cfg)
        self.loss_evaluator = make_roi_box_loss_evaluator(cfg)

    def kl_categorical(self, p_logit, q_logit, weights=None):
        p = F.softmax(p_logit, dim=-1)
        _kl = torch.sum(p * (F.log_softmax(p_logit, dim=-1)
                                  - F.log_softmax(q_logit, dim=-1)), 1)
        if weights is not None:
            _kl *= weights
        return torch.mean(_kl)

    def forward(self, features, proposals, targets=None):
        """
        Arguments:
            features (list[Tensor]): feature-maps from possibly several levels
            proposals (list[BoxList]): proposal boxes
            targets (list[BoxList], optional): the ground-truth targets.

        Returns:
            x (Tensor): the result of the feature extractor
            proposals (list[BoxList]): during training, the subsampled proposals
                are returned. During testing, the predicted boxlists are returned
            losses (dict[Tensor]): During training, returns the losses for the
                head. During testing, returns an empty dict.
        """
        if self.training:
            # 1) 先做正负样本采样 —— 这一步会在 loss_evaluator 内部写入 
            # self._sampled_proposals / self._labels / self._regression_targets
            proposals = self.loss_evaluator.subsample(proposals, targets)
            # 2) 基于“采样后的 proposals”做 RoI 特征抽取
            x = self.feature_extractor(features, proposals)
            # 3) 分类和回归预测
            class_logits, box_regression = self.predictor(x)
            # 4) 计算分类 & 回归的损失（注意这里传 Tensor，不要传 list）
            loss_classifier, loss_box_reg = self.loss_evaluator(class_logits, box_regression)
            # 训练阶段不需要返回 detections，按 maskrcnn_benchmark 的约定返回 (x, proposals, losses)
            return x, proposals, dict(loss_classifier=loss_classifier, loss_box_reg=loss_box_reg)
        else:
            # eval 时不做采样，直接用 RPN 的 proposals
            x = self.feature_extractor(features, proposals)
            class_logits, box_regression = self.predictor(x)
            result = self.post_processor((class_logits, box_regression), proposals)
            return x, result, {}


def build_roi_box_head(cfg, in_channels):
    """
    Constructs a new box head.
    By default, uses ROIBoxHead, but if it turns out not to be enough, just register a new class
    and make it a parameter in the config
    """
    return ROIBoxHead(cfg, in_channels)